{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "38ce0ef6-702b-454c-904b-50d0687796b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\ProgramData\\miniconda3\\envs\\deepseek\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "from glob import glob\n",
    "import os\n",
    "import chromadb\n",
    "\n",
    "\n",
    "os.chdir('D:/AI全栈')\n",
    "\n",
    "#初始化嵌入模型\n",
    "embedding_model = SentenceTransformer(\"shibing624/text2vec-base-chinese\")\n",
    "\n",
    "#把知识向量化\n",
    "def embed_chunk(chunk: str)->list[float]:\n",
    "    embedding = embedding_model.encode(chunk)\n",
    "    return embedding.tolist()\n",
    "\n",
    "#打开知识文件，读取知识数组\n",
    "with open(\"study_ai/dify/无人机知识.md\", \"r\", encoding='utf-8') as file:\n",
    "    file_text = file.read()\n",
    "chunks = file_text.split(\"\\n\\n\")\n",
    "\n",
    "#生成知识的向量化数组\n",
    "embeddings = [embed_chunk(chunk) for chunk in chunks]\n",
    "\n",
    "#打开向量数据库\n",
    "chromadb_client = chromadb.PersistentClient(\"D:/AI全栈/rag/choroma.db\")\n",
    "#打开向量表格\n",
    "chromadb_collection = chromadb_client.get_or_create_collection(name=\"default\")\n",
    "\n",
    "#保存向量到向量库中表格\n",
    "def save_embeddings(chunks: list[str], embeddings: list[list[float]])->None:\n",
    "    ids = [str(i) for i in range(len(chunks))]\n",
    "    chromadb_collection.add(\n",
    "        documents = chunks,\n",
    "        embeddings = embeddings,\n",
    "        ids = ids\n",
    "    )\n",
    "\n",
    "#把知识数据和知识向量存储到向量数据库\n",
    "save_embeddings(chunks, embeddings)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f975ed0b-6115-4a92-a473-0e71de7af232",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0] 问：目前最好的消费级无人机品牌是什么？\n",
      "答：大疆(DJI)在消费级市场占据领先地位\n",
      "\n",
      "[1] 问：无人机飞行记录如何查看？\n",
      "答：通过配套APP可查看详细飞行日志和轨迹回放\n",
      "\n",
      "[2] 问：无人机拍摄需要申请许可吗？\n",
      "答：商业拍摄需申请，个人非营利性使用视当地法规而定\n",
      "\n",
      "[3] 问：农业植保无人机哪个品牌专业？\n",
      "答：极飞(XAG)载重量达80公斤，支持智能播撒和变量喷洒\n",
      "\n",
      "[4] 问：无人机飞行场地怎么选择？\n",
      "答：开阔无遮挡，远离人群、机场和高压线\n",
      "\n",
      "[0] 问：目前最好的消费级无人机品牌是什么？\n",
      "答：大疆(DJI)在消费级市场占据领先地位\n",
      "\n",
      "[1] 问：农业植保无人机哪个品牌专业？\n",
      "答：极飞(XAG)载重量达80公斤，支持智能播撒和变量喷洒\n",
      "\n",
      "[2] 问：无人机拍摄需要申请许可吗？\n",
      "答：商业拍摄需申请，个人非营利性使用视当地法规而定\n",
      "\n",
      "------------------------------\n",
      "\n",
      "问：目前最好的消费级无人机品牌是什么？\n",
      "答：大疆(DJI)在消费级市场占据领先地位\n",
      "问：农业植保无人机哪个品牌专业？\n",
      "答：极飞(XAG)载重量达80公斤，支持智能播撒和变量喷洒\n",
      "问：无人机拍摄需要申请许可吗？\n",
      "答：商业拍摄需申请，个人非营利性使用视当地法规而定\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "from sentence_transformers import CrossEncoder\n",
    "import ollama\n",
    "import chromadb\n",
    "\n",
    "#初始化嵌入模型\n",
    "embedding_model = SentenceTransformer(\"shibing624/text2vec-base-chinese\")\n",
    "\n",
    "#打开向量数据库\n",
    "chromadb_client = chromadb.PersistentClient(\"D:/AI全栈/rag/choroma.db\")\n",
    "#打开向量表格\n",
    "chromadb_collection = chromadb_client.get_or_create_collection(name=\"default\")\n",
    "\n",
    "#把知识向量化\n",
    "def embed_chunk(chunk: str)->list[float]:\n",
    "    embedding = embedding_model.encode(chunk)\n",
    "    return embedding.tolist()\n",
    "\n",
    "#从问题召回top-k个向量\n",
    "def retrieve(query: str, top_k: int)->list[str]:\n",
    "    query_embedding = embed_chunk(query)\n",
    "    results = chromadb_collection.query(\n",
    "        query_embeddings = [query_embedding],\n",
    "        n_results = top_k\n",
    "    )\n",
    "    return results['documents'][0]\n",
    "\n",
    "query = \"无人机品牌有哪些\"\n",
    "retrieved_chunks = retrieve(query, 5)\n",
    "\n",
    "for i, chunk in enumerate(retrieved_chunks):\n",
    "    print(f\"[{i}] {chunk}\\n\")\n",
    "\n",
    "#重排过程\n",
    "def rerank(query: str, retrieved_chunks: list[str], top_k: int)->list[str]:\n",
    "    cross_encoder = CrossEncoder('cross-encoder/mmarco-mMiniLMv2-L12-H384-v1')\n",
    "    pairs = [(query, chunk) for chunk in retrieved_chunks]\n",
    "    scores = cross_encoder.predict(pairs)\n",
    "    chunk_with_score_list = [(chunk, score) for chunk, score in zip(retrieved_chunks, scores)]\n",
    "    chunk_with_score_list.sort(key=lambda pair: pair[1], reverse=True)\n",
    "    return [chunk for chunk, _ in chunk_with_score_list][:top_k]\n",
    "\n",
    "reranked_chunks = rerank(query, retrieved_chunks, 3)\n",
    "for i, chunk in enumerate(reranked_chunks):\n",
    "    print(f\"[{i}] {chunk}\\n\")\n",
    "\n",
    "\n",
    "context = \"\\n\".join(\n",
    "    [line_with_distance for line_with_distance in reranked_chunks]\n",
    ")\n",
    "\n",
    "print(\"------------------------------\\n\")\n",
    "print(f\"{context}\\n\")\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "313ab7ec-c778-4248-bf9c-516a10e2a56f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----------------------------------------\n",
      "<translated>DJi、XAG</translated>\n",
      "\n",
      "-----------------------------------------\n",
      "总耗时：2688902900\n",
      "-----------------------------------------\n"
     ]
    }
   ],
   "source": [
    "SYSTEM_PROMPT = \"\"\"\n",
    "Human: 你是一个 AI 助手。你能够从提供的上下文段落片段中找到问题的答案。\n",
    "\"\"\"\n",
    "\n",
    "USER_PROMPT = f\"\"\"\n",
    "Human: 你是一个 AI 助手。你能够从提供的上下文段落片段中找到问题的答案。\n",
    "请使用以下用 <context> 标签括起来的信息片段来回答用 <question> 标签括起来的问题。最后追加原始回答的中文翻译，并用 <translated>和</translated> 标签标注。\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "<question>\n",
    "{query}\n",
    "</question>\n",
    "<translated>\n",
    "</translated>\n",
    "\"\"\"\n",
    "\n",
    "stream = ollama.chat(\n",
    "    stream=True,\n",
    "    model='llama3.2:latest', # 修改大模型名称1\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
    "        {\"role\": \"user\", \"content\": USER_PROMPT}\n",
    "    ]\n",
    ")\n",
    "\n",
    "print('-----------------------------------------')\n",
    "for chunk in stream:\n",
    "    if not chunk['done']:\n",
    "        print(chunk['message']['content'], end=\"\", flush=True)\n",
    "    else:\n",
    "        print('\\n')\n",
    "        print('-----------------------------------------')\n",
    "        print(f'总耗时：{chunk['total_duration']}')\n",
    "        print('-----------------------------------------')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f463709c-eebc-4768-98f6-7a21af77e0d1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
